import torch
import torch.nn as nn


def get_minimal_tensor_info(indices, B, D, H, W):
    min_b = min([b for b, _, _, _ in indices])
    max_b = max([b for b, _, _, _ in indices])
    min_d = min([d for _, d, _, _ in indices])
    max_d = max([d for _, d, _, _ in indices])
    min_h = min([h for _, _, h, _ in indices])
    max_h = max([h for _, _, h, _ in indices])
    min_w = min([w for _, _, _, w in indices])
    max_w = max([w for _, _, _, w in indices])

    new_b = max_b - min_b + 1
    new_d = max_d - min_d + 1
    new_h = max_h - min_h + 1
    new_w = max_w - min_w + 1

    return min_b, max_b, min_d, max_d, min_h, max_h, min_w, max_w, new_b, new_d, new_h, new_w


def create_mapping(indices, min_b, min_d, min_h, min_w):
    mapping = {}
    for i, (b, d, h, w) in enumerate(indices):
        new_b = b - min_b
        new_d = d - min_d
        new_h = h - min_h
        new_w = w - min_w
        mapping[i] = (new_b, new_d, new_h, new_w)
    return mapping


def fill_minimal_tensor(minimal_tensor, features, mapping):
    for i, feature in enumerate(features):
        new_b, new_d, new_h, new_w = mapping[i]
        minimal_tensor[new_b, new_d, new_h, new_w] = feature
    return minimal_tensor


def recover_features(output, mapping, min_b, min_d, min_h, min_w, B, D, H, W, C):
    original_output = torch.zeros((B, D, H, W, C), device=output.device)
    for i, (new_b, new_d, new_h, new_w) in mapping.items():
        b = min_b + new_b
        d = min_d + new_d
        h = min_h + new_h
        w = min_w + new_w
        original_output[b, d, h, w] = output[new_b, new_d, new_h, new_w]
    return original_output


# 示例使用
B, D, H, W, C = 2, 3, 4, 5, 6
tensor = torch.zeros((B, D, H, W, C))
indices = torch.Tensor([[0, 1, 2, 3], [1, 2, 3, 4]])
features = torch.Tensor([[0, 1, 2, 3], [1, 2, 3, 4]])

min_b, max_b, min_d, max_d, min_h, max_h, min_w, max_w, new_b, new_d, new_h, new_w = get_minimal_tensor_info(indices, B, D, H, W)
mapping = create_mapping(indices, min_b, min_d, min_h, min_w)

minimal_tensor = torch.zeros((new_b, new_d, new_h, new_w, C))
filled_minimal_tensor = fill_minimal_tensor(minimal_tensor, features, mapping)

# 进行3D卷积计算
conv3d = nn.Conv3d(C, C, kernel_size=3, stride=1, padding=1)
output = conv3d(filled_minimal_tensor.permute(4, 0, 1, 2, 3)).permute(1, 2, 3, 4, 0)

# 恢复原始特征顺序
original_output = recover_features(output, mapping, min_b, min_d, min_h, min_w, B, D, H, W, C)

print(original_output.shape)